import os
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.cuda.amp import autocast, GradScaler
from tqdm import tqdm
import shutil
import logging
import json
from os import path
from utils import HParams

logging.getLogger('numba').setLevel(logging.WARNING)

import commons
import utils
from data_utils import (
    TextAudioSpeakerLoader,
    TextAudioSpeakerCollate,
    DistributedBucketSampler
)
from models import (
    SynthesizerTrn,
    MultiPeriodDiscriminator,
)
from losses import (
    generator_loss,
    discriminator_loss,
    feature_loss,
    kl_loss
)
from mel_processing import mel_spectrogram_torch, spec_to_mel_torch

torch.backends.cudnn.benchmark = True
global_step = 0
current_path = os.path.dirname(os.path.abspath(__file__))




def run(rank, n_gpus, hps, hps_dict, pre_trained_G_path):
    global global_step
    symbols = hps['symbols']
    dist.init_process_group(backend='gloo' if os.name == 'nt' else 'nccl', init_method='env://', world_size=n_gpus,
                            rank=rank)
    torch.manual_seed(hps.train.seed)
    torch.cuda.set_device(rank)

    train_dataset = TextAudioSpeakerLoader(hps.data.training_files, hps.data, symbols)
    train_sampler = DistributedBucketSampler(
        train_dataset,
        hps.train.batch_size,
        [32, 300, 400, 500, 600, 700, 800, 900, 1000],
        num_replicas=n_gpus,
        rank=rank,
        shuffle=True)
    collate_fn = TextAudioSpeakerCollate()
    train_loader = DataLoader(train_dataset, num_workers=2, shuffle=False, pin_memory=True,
                              collate_fn=collate_fn, batch_sampler=train_sampler)
    # train_loader = DataLoader(train_dataset, batch_size=hps.train.batch_size, num_workers=2, shuffle=False, pin_memory=True,
    #                           collate_fn=collate_fn)

    net_g = SynthesizerTrn(
        len(symbols),
        hps.data.filter_length // 2 + 1,
        hps.train.segment_size // hps.data.hop_length,
        n_speakers=hps.data.n_speakers,
        **hps.model).cuda(rank)
    net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(rank)
    # freeze all other layers except speaker embedding
    for p in net_g.parameters():
        p.requires_grad = True
    for p in net_d.parameters():
        p.requires_grad = True
    # for p in net_d.parameters():
    #     p.requires_grad = False
    # net_g.emb_g.weight.requires_grad = True
    optim_g = torch.optim.AdamW(
        net_g.parameters(),
        hps.train.learning_rate,
        betas=hps.train.betas,
        eps=hps.train.eps)
    optim_d = torch.optim.AdamW(
        net_d.parameters(),
        hps.train.learning_rate,
        betas=hps.train.betas,
        eps=hps.train.eps)
    # optim_d = None
    net_g = DDP(net_g, device_ids=[rank])
    net_d = DDP(net_d, device_ids=[rank])
    # load existing model
    if pre_trained_G_path:
        _, _, _, epoch_str = utils.load_checkpoint(pre_trained_G_path, net_g, None,
                                                   drop_speaker_emb=hps.drop_speaker_embed)

        _, _, _, _ = utils.load_checkpoint(path.join(current_path, "pretrained_models/D_0.pth"), net_d, optim_d)
    else:
        _, _, _, _ = utils.load_checkpoint(path.join(current_path, "pretrained_models/G_0.pth"), net_g, None,
                                           drop_speaker_emb=hps.drop_speaker_embed)

        _, _, _, _ = utils.load_checkpoint(path.join(current_path, "pretrained_models/D_0.pth"), net_d, optim_d)
        epoch_str = 0

    global_step = epoch_str * len(train_loader) + 1

    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay)
    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay)

    scaler = GradScaler(enabled=hps.train.fp16_run)

    for epoch in range(epoch_str + 1, hps.train.epochs + 1):
        if rank == 0:
            train_and_evaluate(rank, epoch, hps, hps_dict, [net_g, net_d], [optim_g, optim_d],
                               [scheduler_g, scheduler_d], scaler,
                               [train_loader, None], None, None)
        else:
            train_and_evaluate(rank, epoch, hps, hps_dict, [net_g, net_d], [optim_g, optim_d],
                               [scheduler_g, scheduler_d], scaler,
                               [train_loader, None], None, None)
        scheduler_g.step()
        scheduler_d.step()


def train_and_evaluate(rank, epoch, hps, hps_dict, nets, optims, schedulers, scaler, loaders, logger, writers):
    net_g, net_d = nets
    optim_g, optim_d = optims
    scheduler_g, scheduler_d = schedulers
    train_loader, eval_loader = loaders
    if writers is not None:
        writer, writer_eval = writers

    # train_loader.batch_sampler.set_epoch(epoch)
    global global_step

    net_g.train()
    net_d.train()
    for batch_idx, (x, x_lengths, spec, spec_lengths, y, y_lengths, speakers) in enumerate(tqdm(train_loader)):
        x, x_lengths = x.cuda(rank, non_blocking=True), x_lengths.cuda(rank, non_blocking=True)
        spec, spec_lengths = spec.cuda(rank, non_blocking=True), spec_lengths.cuda(rank, non_blocking=True)
        y, y_lengths = y.cuda(rank, non_blocking=True), y_lengths.cuda(rank, non_blocking=True)
        speakers = speakers.cuda(rank, non_blocking=True)

        with autocast(enabled=hps.train.fp16_run):
            y_hat, l_length, attn, ids_slice, x_mask, z_mask, \
            (z, z_p, m_p, logs_p, m_q, logs_q) = net_g(x, x_lengths, spec, spec_lengths, speakers)

            mel = spec_to_mel_torch(
                spec,
                hps.data.filter_length,
                hps.data.n_mel_channels,
                hps.data.sampling_rate,
                hps.data.mel_fmin,
                hps.data.mel_fmax)
            y_mel = commons.slice_segments(mel, ids_slice, hps.train.segment_size // hps.data.hop_length)
            y_hat_mel = mel_spectrogram_torch(
                y_hat.squeeze(1),
                hps.data.filter_length,
                hps.data.n_mel_channels,
                hps.data.sampling_rate,
                hps.data.hop_length,
                hps.data.win_length,
                hps.data.mel_fmin,
                hps.data.mel_fmax
            )

            y = commons.slice_segments(y, ids_slice * hps.data.hop_length, hps.train.segment_size)  # slice

            # Discriminator
            y_d_hat_r, y_d_hat_g, _, _ = net_d(y, y_hat.detach())
            with autocast(enabled=False):
                loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
                loss_disc_all = loss_disc
        optim_d.zero_grad()
        scaler.scale(loss_disc_all).backward()
        scaler.unscale_(optim_d)
        grad_norm_d = commons.clip_grad_value_(net_d.parameters(), None)
        scaler.step(optim_d)

        with autocast(enabled=hps.train.fp16_run):
            # Generator
            y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y, y_hat)
            with autocast(enabled=False):
                loss_dur = torch.sum(l_length.float())
                loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel
                loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.c_kl

                loss_fm = feature_loss(fmap_r, fmap_g)
                loss_gen, losses_gen = generator_loss(y_d_hat_g)
                loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl
        optim_g.zero_grad()
        scaler.scale(loss_gen_all).backward()
        scaler.unscale_(optim_g)
        grad_norm_g = commons.clip_grad_value_(net_g.parameters(), None)
        scaler.step(optim_g)
        scaler.update()
        if rank == 0:
            if global_step % hps.train.eval_interval == 0:
                model_save_path = os.path.join(hps.model_dir, f'G_{global_step}/G_{global_step}.pth')
                if not os.path.exists(os.path.dirname(model_save_path)):
                    os.makedirs(os.path.dirname(model_save_path))
                # 保存配置文件
                with open(os.path.join(os.path.dirname(model_save_path), 'config.json'), 'w', encoding='utf-8') as f:
                    json.dump(hps_dict, f, indent=2)
                utils.save_checkpoint(net_g, None, hps.train.learning_rate, epoch, model_save_path)
                old_g = os.path.join(hps.model_dir, "G_{}".format(global_step - 38000))
                if os.path.exists(old_g):
                    shutil.rmtree(old_g)
        global_step += 1
    if rank == 0:
        print(f'====> Epoch: {epoch} finish!\n')
        if hps.train.epochs == epoch:
            model_save_path = os.path.join(hps.model_dir, f'G_{global_step - 1}/G_{global_step - 1}.pth')
            if not os.path.exists(os.path.dirname(model_save_path)):
                os.makedirs(os.path.dirname(model_save_path))
            # 保存配置文件
            with open(os.path.join(os.path.dirname(model_save_path), 'config.json'), 'w', encoding='utf-8') as f:
                json.dump(hps_dict, f, indent=2)
            utils.save_checkpoint(net_g, None, hps.train.learning_rate, epoch, model_save_path)

def main(hps_dict, pre_trained_G_path=None):
    """Assume Single Node Multi GPUs Training Only"""
    assert torch.cuda.is_available(), "CPU training is not allowed."
    n_gpus = torch.cuda.device_count()
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '8000'
    hps = HParams(**hps_dict)
    # run(0, n_gpus, hps, hps_dict, pre_trained_G_path)
    mp.spawn(run, nprocs=n_gpus, args=(n_gpus, hps, hps_dict, pre_trained_G_path))



if __name__ == "__main__":
    main()
